
#include "StdAfx.h"
#include <shlobj.h>
#include <atlbase.h>
#include "DragDropImpl.h"

// CEnumFormatEtc
CEnumFormatEtc::CEnumFormatEtc(const CSimpleArray<FORMATETC>& ArrFE)
: m_cRef(0), m_iCur(0)
{
	for (int i = 0; i < ArrFE.GetSize(); ++i)
		m_pFmtEtc.Add(ArrFE[i]);
}

CEnumFormatEtc::CEnumFormatEtc(const CSimpleArray<FORMATETC*>& ArrFE)
: m_cRef(0), m_iCur(0)
{
	for (int i = 0; i < ArrFE.GetSize(); ++i)
		m_pFmtEtc.Add(*ArrFE[i]);
}

STDMETHODIMP CEnumFormatEtc::QueryInterface(REFIID riid, void** ppvObject)
{
	if (ppvObject)
		return E_INVALIDARG;

	*ppvObject = NULL;

	if (IID_IUnknown == riid || IID_IEnumFORMATETC == riid)
		*ppvObject = this;

	if (*ppvObject != NULL)
	{
		((LPUNKNOWN)*ppvObject)->AddRef();
		return S_OK;
	}

	return E_NOINTERFACE;
}

STDMETHODIMP_(ULONG) CEnumFormatEtc::AddRef()
{
	return ++m_cRef;
}

STDMETHODIMP_(ULONG) CEnumFormatEtc::Release()
{
	long nTemp = --m_cRef;
	ATLASSERT(nTemp >= 0);

	if(nTemp == 0)
		delete this;

	return nTemp; 
}

STDMETHODIMP CEnumFormatEtc::Next(ULONG celt, LPFORMATETC lpFormatEtc, ULONG *pceltFetched)
{
	if (pceltFetched != NULL)
		*pceltFetched=0;

	ULONG cReturn = celt;

	if (celt <= 0 || lpFormatEtc == NULL || m_iCur >= m_pFmtEtc.GetSize())
		return S_FALSE;

	if (pceltFetched == NULL && celt != 1)
		return S_FALSE;

	while (m_iCur < m_pFmtEtc.GetSize() && cReturn > 0)
	{
		*lpFormatEtc++ = m_pFmtEtc[m_iCur++];
		--cReturn;
	}

	if (pceltFetched != NULL)
		*pceltFetched = celt - cReturn;

	return (cReturn == 0) ? S_OK : S_FALSE;
}

STDMETHODIMP CEnumFormatEtc::Skip(ULONG celt)
{
	if ((m_iCur + int(celt)) >= m_pFmtEtc.GetSize())
		return S_FALSE;

	m_iCur += celt;
	return S_OK;
}

STDMETHODIMP CEnumFormatEtc::Reset(void)
{
	m_iCur = 0;
	return S_OK;
}

STDMETHODIMP CEnumFormatEtc::Clone(IEnumFORMATETC **ppCloneEnumFormatEtc)
{
	if (ppCloneEnumFormatEtc == NULL)
		return E_POINTER;

	CEnumFormatEtc *newEnum = new CEnumFormatEtc(m_pFmtEtc);
	if(newEnum ==NULL)
		return E_OUTOFMEMORY;

	newEnum->AddRef();
	newEnum->m_iCur = m_iCur;
	*ppCloneEnumFormatEtc = newEnum;

	return S_OK;
}

// CIDropSource Class
CIDropSource::CIDropSource() : m_cRef(0)
{
}

STDMETHODIMP CIDropSource::QueryInterface(REFIID riid, void **ppvObject)
{
	if (ppvObject == NULL)
		return E_INVALIDARG;

	*ppvObject = NULL;

	if (IID_IUnknown == riid || IID_IDropSource == riid)
		*ppvObject = this;

	if (*ppvObject != NULL)
	{
		((LPUNKNOWN)*ppvObject)->AddRef();
		return S_OK;
	}

	return E_NOINTERFACE;
}

STDMETHODIMP_(ULONG) CIDropSource::AddRef()
{
	return ++m_cRef;
}

STDMETHODIMP_(ULONG) CIDropSource::Release()
{
	long nTemp = --m_cRef;
	ATLASSERT(nTemp >= 0);

	if (nTemp == 0)
		delete this;

	return nTemp;
}

STDMETHODIMP CIDropSource::QueryContinueDrag(BOOL fEscapePressed, DWORD grfKeyState)
{
	if (fEscapePressed)
		return DRAGDROP_S_CANCEL;

	if (!(grfKeyState & (MK_LBUTTON | MK_RBUTTON)))
		return DRAGDROP_S_DROP;

	return S_OK;

}

STDMETHODIMP CIDropSource::GiveFeedback(DWORD dwEffect)
{
	return DRAGDROP_S_USEDEFAULTCURSORS;
	//return S_OK;
}

// CIDataObject
CIDataObject::CIDataObject(CIDropSource* pDropSource)
: m_cRef(0), m_pDropSource(pDropSource)
{
}

CIDataObject::~CIDataObject()
{
	for (int i = 0; i < m_StgMedium.GetSize(); ++i)
	{
		ReleaseStgMedium(m_StgMedium[i]);
		delete m_StgMedium[i];
	}

	for (int j = 0; j < m_ArrFormatEtc.GetSize(); ++j)
		delete m_ArrFormatEtc[j];
}

STDMETHODIMP CIDataObject::QueryInterface(REFIID riid, void **ppvObject)
{
	if (ppvObject == NULL)
		return E_INVALIDARG;

	*ppvObject = NULL;

	if (IID_IUnknown == riid || IID_IDataObject == riid)
		*ppvObject=this;

	if (NULL != *ppvObject)
	{
		((LPUNKNOWN)*ppvObject)->AddRef();
		return S_OK;
	}

	return E_NOINTERFACE;
}

STDMETHODIMP_(ULONG) CIDataObject::AddRef()
{
	return ++m_cRef;
}

STDMETHODIMP_(ULONG) CIDataObject::Release()
{
	long nTemp = --m_cRef;
	if (nTemp == 0)
		delete this;

	return nTemp;
}

STDMETHODIMP CIDataObject::GetData(FORMATETC *pformatetcIn, STGMEDIUM *pmedium)
{ 
	if (pformatetcIn == NULL || pmedium == NULL)
		return E_INVALIDARG;

	pmedium->hGlobal = NULL;

	ATLASSERT(m_StgMedium.GetSize() == m_ArrFormatEtc.GetSize());
	for (int i = 0; i < m_ArrFormatEtc.GetSize(); ++i)
	{
		if (pformatetcIn->tymed & m_ArrFormatEtc[i]->tymed &&
			pformatetcIn->dwAspect == m_ArrFormatEtc[i]->dwAspect &&
			pformatetcIn->cfFormat == m_ArrFormatEtc[i]->cfFormat)
		{
			CopyMedium(pmedium, m_StgMedium[i], m_ArrFormatEtc[i]);
			return S_OK;
		}
	}

	return DV_E_FORMATETC;
}

STDMETHODIMP CIDataObject::GetDataHere(FORMATETC __RPC_FAR *pformatetc, STGMEDIUM *pmedium)
{ 
	return E_NOTIMPL;
}

STDMETHODIMP CIDataObject::QueryGetData(FORMATETC *pformatetc)
{ 
	if (pformatetc == NULL)
		return E_INVALIDARG;

	if (!(DVASPECT_CONTENT & pformatetc->dwAspect))
		return DV_E_DVASPECT;

	HRESULT hr = DV_E_TYMED;
	for (int i = 0; i < m_ArrFormatEtc.GetSize(); ++i)
	{
		if (pformatetc->tymed & m_ArrFormatEtc[i]->tymed)
		{
			if (pformatetc->cfFormat == m_ArrFormatEtc[i]->cfFormat)
				return S_OK;
			else
				hr = DV_E_CLIPFORMAT;
		}
		else
			hr = DV_E_TYMED;
	}

	return hr;
}

STDMETHODIMP CIDataObject::GetCanonicalFormatEtc(FORMATETC *pformatectIn, FORMATETC *pformatetcOut)
{ 
	if (pformatetcOut == NULL)
		return E_INVALIDARG;

	return DATA_S_SAMEFORMATETC;
}

STDMETHODIMP CIDataObject::SetData(FORMATETC *pformatetc, STGMEDIUM *pmedium, BOOL fRelease)
{ 
	if (pformatetc == NULL || pmedium == NULL)
		return E_INVALIDARG;

	ATLASSERT(pformatetc->tymed == pmedium->tymed);
	FORMATETC* fetc = new FORMATETC;
	STGMEDIUM* pStgMed = new STGMEDIUM;

	if (fetc == NULL || pStgMed == NULL)
		return E_OUTOFMEMORY;

	ZeroMemory(fetc,sizeof(FORMATETC));
	ZeroMemory(pStgMed,sizeof(STGMEDIUM));

	*fetc = *pformatetc;
	m_ArrFormatEtc.Add(fetc);

	if (fRelease)
		*pStgMed = *pmedium;
	else
	{
		CopyMedium(pStgMed, pmedium, pformatetc);
	}

	m_StgMedium.Add(pStgMed);

	return S_OK;
}
void CIDataObject::CopyMedium(STGMEDIUM* pMedDest, STGMEDIUM* pMedSrc, FORMATETC* pFmtSrc)
{
	switch(pMedSrc->tymed)
	{
	case TYMED_HGLOBAL:
		pMedDest->hGlobal = (HGLOBAL)OleDuplicateData(pMedSrc->hGlobal,pFmtSrc->cfFormat, NULL);
		break;

	case TYMED_GDI:
		pMedDest->hBitmap = (HBITMAP)OleDuplicateData(pMedSrc->hBitmap,pFmtSrc->cfFormat, NULL);
		break;

	case TYMED_MFPICT:
		pMedDest->hMetaFilePict = (HMETAFILEPICT)OleDuplicateData(pMedSrc->hMetaFilePict,pFmtSrc->cfFormat, NULL);
		break;

	case TYMED_ENHMF:
		pMedDest->hEnhMetaFile = (HENHMETAFILE)OleDuplicateData(pMedSrc->hEnhMetaFile,pFmtSrc->cfFormat, NULL);
		break;

	case TYMED_FILE:
		pMedSrc->lpszFileName = (LPOLESTR)OleDuplicateData(pMedSrc->lpszFileName,pFmtSrc->cfFormat, NULL);
		break;

	case TYMED_ISTREAM:
		pMedDest->pstm = pMedSrc->pstm;
		pMedSrc->pstm->AddRef();
		break;

	case TYMED_ISTORAGE:
		pMedDest->pstg = pMedSrc->pstg;
		pMedSrc->pstg->AddRef();
		break;

	case TYMED_NULL:
	default:
		break;
	}

	pMedDest->tymed = pMedSrc->tymed;
	pMedDest->pUnkForRelease = NULL;
	if (pMedSrc->pUnkForRelease != NULL)
	{
		pMedDest->pUnkForRelease = pMedSrc->pUnkForRelease;
		pMedSrc->pUnkForRelease->AddRef();
	}
}

STDMETHODIMP CIDataObject::EnumFormatEtc(DWORD dwDirection, IEnumFORMATETC **ppenumFormatEtc)
{ 
	if (ppenumFormatEtc == NULL)
		return E_POINTER;

	*ppenumFormatEtc=NULL;
	switch (dwDirection)
	{
	case DATADIR_GET:
		*ppenumFormatEtc= new CEnumFormatEtc(m_ArrFormatEtc);
		if (*ppenumFormatEtc == NULL)
			return E_OUTOFMEMORY;
		(*ppenumFormatEtc)->AddRef(); 
		break;

	case DATADIR_SET:
	default:
		return E_NOTIMPL;
		break;
	}

	return S_OK;
}

STDMETHODIMP CIDataObject::DAdvise(FORMATETC *pformatetc, DWORD advf, IAdviseSink *pAdvSink, DWORD *pdwConnection)
{ 
	return OLE_E_ADVISENOTSUPPORTED;
}

STDMETHODIMP CIDataObject::DUnadvise(DWORD dwConnection)
{
	return E_NOTIMPL;
}

HRESULT STDMETHODCALLTYPE CIDataObject::EnumDAdvise(IEnumSTATDATA **ppenumAdvise)
{
	return OLE_E_ADVISENOTSUPPORTED;
}

// CIDropTarget Class
CIDropTarget::CIDropTarget(bool bSelfDelete)
: m_bSelfDelete(bSelfDelete), m_cRef(0), m_bAllowDrop(false), m_pDropTargetHelper(NULL), m_pSupportedFrmt(NULL)
{ 
}

CIDropTarget::~CIDropTarget()
{
	if(m_pDropTargetHelper != NULL)
	{
		m_pDropTargetHelper->Release();
		m_pDropTargetHelper = NULL;
	}
}

STDMETHODIMP CIDropTarget::QueryInterface(REFIID riid, void **ppvObject)
{
	if (ppvObject == NULL)
		return E_INVALIDARG;

	*ppvObject = NULL;

	if (IID_IUnknown == riid || IID_IDropTarget == riid)
		*ppvObject = this;

	if (*ppvObject != NULL)
	{
		((LPUNKNOWN)*ppvObject)->AddRef();
		return S_OK;
	}

	return E_NOINTERFACE;
}

STDMETHODIMP_(ULONG) CIDropTarget::AddRef()
{ 
	return ++m_cRef;
}

ULONG STDMETHODCALLTYPE CIDropTarget::Release()
{
	long nTemp = --m_cRef;
	ATLASSERT(nTemp >= 0);

	if (nTemp == 0 && m_bSelfDelete)
		delete this;

	return nTemp;
}

STDMETHODIMP CIDropTarget::DragEnter(IDataObject *pDataObj, DWORD grfKeyState, POINTL pt, DWORD *pdwEffect)
{
	if (pDataObj == NULL)
		return E_INVALIDARG;

	if (m_pDropTargetHelper)
		m_pDropTargetHelper->DragEnter(m_hWndTarget, pDataObj, (LPPOINT)&pt, *pdwEffect);

	m_pSupportedFrmt = NULL;
	for (int i =0; i< m_arrFormatEtc.GetSize(); ++i)
	{
		m_bAllowDrop = (pDataObj->QueryGetData(&m_arrFormatEtc[i]) == S_OK);
		if (m_bAllowDrop)
		{
			m_pSupportedFrmt = &m_arrFormatEtc[i];
			break;
		}
	}

	QueryDrop(grfKeyState, pdwEffect);

	return S_OK;
}

STDMETHODIMP CIDropTarget::DragOver(DWORD grfKeyState, POINTL pt, DWORD *pdwEffect)
{
	if (m_pDropTargetHelper)
		m_pDropTargetHelper->DragOver((LPPOINT)&pt, *pdwEffect);
	
	QueryDrop(grfKeyState, pdwEffect);
	
	return S_OK;
}

STDMETHODIMP CIDropTarget::DragLeave()
{
	if (m_pDropTargetHelper)
		m_pDropTargetHelper->DragLeave();

	m_bAllowDrop = false;
	m_pSupportedFrmt = NULL;

	return S_OK;
}

STDMETHODIMP CIDropTarget::Drop(IDataObject *pDataObj, DWORD grfKeyState, POINTL pt, DWORD *pdwEffect)
{
	if (pDataObj == NULL)
		return E_INVALIDARG;	

	if (m_pDropTargetHelper)
		m_pDropTargetHelper->Drop(pDataObj, (LPPOINT)&pt, *pdwEffect);

	std::vector<int> vecIndices;
	if (QueryDrop(grfKeyState, pdwEffect))
	{
		if (m_bAllowDrop && m_pSupportedFrmt != NULL)
		{
			STGMEDIUM medium;
			if (pDataObj->GetData(m_pSupportedFrmt, &medium) == S_OK)
			{
				if (OnDrop(m_pSupportedFrmt, medium, pdwEffect))
				{
					ReleaseStgMedium(&medium);
				}
			}
		}
	}

	m_bAllowDrop = false;
	*pdwEffect = DROPEFFECT_NONE;
	m_pSupportedFrmt = NULL;

	return S_OK;
}

bool CIDropTarget::QueryDrop(DWORD grfKeyState, LPDWORD pdwEffect)
{  
	DWORD dwOKEffects = *pdwEffect; 

	if(!m_bAllowDrop)
	{
		*pdwEffect = DROPEFFECT_NONE;
		return false;
	}

	//CTRL+SHIFT  -- DROPEFFECT_LINK
	//CTRL        -- DROPEFFECT_COPY
	//SHIFT       -- DROPEFFECT_MOVE
	//no modifier -- DROPEFFECT_MOVE or whatever is allowed by src
	*pdwEffect = (grfKeyState & MK_CONTROL) ?
		((grfKeyState & MK_SHIFT) ? DROPEFFECT_LINK : DROPEFFECT_COPY):
		((grfKeyState & MK_SHIFT) ? DROPEFFECT_MOVE : DROPEFFECT_MOVE);

	if (*pdwEffect == 0) 
	{
		if (DROPEFFECT_COPY & dwOKEffects)
			*pdwEffect = DROPEFFECT_COPY;
		else if (DROPEFFECT_MOVE & dwOKEffects)
			*pdwEffect = DROPEFFECT_MOVE; 
		else if (DROPEFFECT_LINK & dwOKEffects)
			*pdwEffect = DROPEFFECT_LINK; 
		else 
		{
			*pdwEffect = DROPEFFECT_NONE;
		}
	} 
	else
	{
		if(!(*pdwEffect & dwOKEffects))
			*pdwEffect = DROPEFFECT_NONE;
	}  

	return (DROPEFFECT_NONE != *pdwEffect);
}   

void CIDropTarget::SetTargetWnd(HWND hWnd)
{
	m_hWndTarget = hWnd;
	ATLASSERT(m_hWndTarget != NULL);

	if (FAILED(CoCreateInstance(CLSID_DragDropHelper,NULL,CLSCTX_INPROC_SERVER, 
		IID_IDropTargetHelper, (LPVOID*)&m_pDropTargetHelper)))
	{
		m_pDropTargetHelper = NULL;
	}
}
